{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f1458df-5306-43de-a565-d5f521d0abce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import timeit\n",
    "import math\n",
    "import json\n",
    "import io\n",
    "import pandas as pd\n",
    "import buckaroo\n",
    "import numpy as np\n",
    "from buckaroo.serialization_utils import pd_to_obj\n",
    "\n",
    "N = 500\n",
    "NA = pd.NA\n",
    "ROWS = 100_000\n",
    "typed_df = pd.DataFrame({'int_col':np.random.randint(1,50, ROWS), 'float_col': np.random.randint(1,30, ROWS)/.7,\n",
    "                         \"str_col\": [\"foobar\"]* ROWS,\n",
    "                        'longtail_unique' : random_categorical({},      unique_per=0.5, na_per=.0, longtail_per=.5, N=N),\n",
    "\n",
    "                        })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2828909-0d9c-4628-a1b2-041bbdccbbac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rand_cat(named_p, na_per, N):\n",
    "    choices, p = [], []\n",
    "    named_total_per = sum(named_p.values()) + na_per\n",
    "    total_len = int(np.floor(named_total_per * N))\n",
    "    if named_total_per > 0:\n",
    "        for k, v in named_p.items():\n",
    "            choices.append(k)\n",
    "            p.append(v/named_total_per)\n",
    "\n",
    "        choices.append(pd.NA)\n",
    "        p.append(na_per/named_total_per)    \n",
    "        return [np.random.choice(choices, p=p) for k in range(total_len)]\n",
    "    else:\n",
    "        return []\n",
    "\n",
    "def random_categorical(named_p, unique_per, na_per, longtail_per, N):\n",
    "    choice_arr = rand_cat(named_p, na_per, N)\n",
    "    discrete_choice_len = len(choice_arr)\n",
    "\n",
    "    longtail_count = int(np.floor(longtail_per * N))//2\n",
    "    extra_arr = []\n",
    "    for i in range(longtail_count):\n",
    "        extra_arr.append(\"long_%d\" % i)\n",
    "        extra_arr.append(\"long_%d\" % i)\n",
    "\n",
    "    unique_len = N - (len(extra_arr) + discrete_choice_len)\n",
    "    #print(\"discrete_choice_len\", discrete_choice_len, \"longtail_count\", longtail_count, \"unique_len\", unique_len)\n",
    "    for i in range(unique_len):\n",
    "        extra_arr.append(\"unique_%d\" % i)\n",
    "    all_arr = np.concatenate([choice_arr, extra_arr])\n",
    "    np.random.shuffle(all_arr)\n",
    "    return all_arr        \n",
    "    \n",
    "N = 500\n",
    "NA = pd.NA\n",
    "ROWS = 100_000\n",
    "typed_df = pd.DataFrame({'int_col':np.random.randint(1,50, ROWS), 'float_col': np.random.randint(1,30, ROWS)/.7,\n",
    "                        \"str_col\": [\"foobar\"]* ROWS,\n",
    "                        \"longtail_unique\" : random_categorical({},\n",
    "                            unique_per=0.5, na_per=.0, longtail_per=.5, N=ROWS),\n",
    "                        })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d240e31e-aefd-4fec-9123-7a77b039108f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def human_size_temporal(frac_seconds, order=None):\n",
    "    precision = 3\n",
    "    units = [\"s\", \"ms\",'us',\"ns\"]\n",
    "    scaling = [1, 1e3, 1e6, 1e9]\n",
    "\n",
    "    if order is None:\n",
    "        if frac_seconds > 0.0 and frac_seconds < 1000.0:\n",
    "            order = min(-int(math.floor(math.log10(frac_seconds)) // 3), 3)\n",
    "        elif frac_seconds >= 1000.0:\n",
    "            order = 0\n",
    "        else:\n",
    "            order = 3\n",
    "    #print(\" %.*g %s per loop\" % (precision, frac_seconds * scaling[order],  units[order]))\n",
    "    return (round(frac_seconds*scaling[order], precision), units[order], order)\n",
    "def time_stats(stmt, number=29, repeat=5, order=None, **kwargs):\n",
    "    \"\"\"adapted from ipython magic timeit \n",
    "    https://github.com/jupyter/ipython-py3k/blob/17bd8dfb1da9940799c3c3e0e19af098f4f93a45/IPython/core/magic.py#L1741\n",
    "    \"\"\"\n",
    "    times = timeit.repeat(stmt, number=number, repeat=repeat, globals=kwargs)\n",
    "    time_arr = np.array(times) / number \n",
    "    np.mean(time_arr)\n",
    "    best = np.min(time_arr)\n",
    "\n",
    "    scaled, units, _order = human_size_temporal(best, order=order)\n",
    "    print(\"%d loops, best of %d: %f %s per loop\" % (number, repeat, scaled, units))\n",
    "    return scaled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "856be7d8-96da-4358-9a65-20bcf992e0e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def arrow_write(df):\n",
    "    f = io.BytesIO()\n",
    "    df.to_feather(f)\n",
    "def arrow_size(df):\n",
    "    f = io.BytesIO()\n",
    "    df.to_feather(f)\n",
    "    f.seek(0)\n",
    "    return len(f.read())\n",
    "def run_timings(N, number=29):\n",
    "    ltyped_df=typed_df[:N]\n",
    "    exec_args = dict(N=N, pd_to_obj=pd_to_obj, typed_df=ltyped_df, arrow_write=arrow_write,\n",
    "                     order=1, number=number)\n",
    "    return dict(\n",
    "            N=N,\n",
    "            parquet_time=time_stats('typed_df.to_parquet()', **exec_args),\n",
    "            parquet_size=len(ltyped_df.to_parquet()),\n",
    "            arrow_time=time_stats('arrow_write(typed_df)', **exec_args),\n",
    "            arrow_size=arrow_size(ltyped_df),\n",
    "            pd_to_obj_time=time_stats('pd_to_obj(typed_df)', **exec_args),\n",
    "            pd_to_obj_size=len(json.dumps(pd_to_obj(ltyped_df))))\n",
    "time_df = pd.DataFrame([\n",
    "    run_timings(100),\n",
    "    run_timings(500),\n",
    "    run_timings(1000),\n",
    "    run_timings(5000),\n",
    "    run_timings(10000),\n",
    "    run_timings(20000, 10),\n",
    "    run_timings(100_000, 5),\n",
    "])\n",
    "time_df['parq_speedup'] = time_df['pd_to_obj_time'] / time_df['parquet_time'] \n",
    "time_df['parq_per_byte'] = (time_df['parquet_time'] / time_df['parquet_size']) * 1_000_000\n",
    "time_df['json_per_byte'] = (time_df['pd_to_obj_time'] / time_df['pd_to_obj_size']) * 1_000_000\n",
    "buckaroo.BuckarooWidget(time_df)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d709c036-a6ed-430c-bdc5-3dbc7c457564",
   "metadata": {},
   "outputs": [],
   "source": [
    "typed_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25b80d05-3ca3-4ebc-ab1b-fd9ea62cc07c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
